package com.projects.modular.system.service;

import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Service
public class SseService {

    private final Map<String, SseEmitter> emitters = new ConcurrentHashMap<>();

    /**
     * 添加客户端连接
     */
    public void addEmitter(String clientId, SseEmitter emitter) {
        emitters.put(clientId, emitter);
        
        // 注册清理回调
        emitter.onCompletion(() -> emitters.remove(clientId));
        emitter.onTimeout(() -> emitters.remove(clientId));
    }

    /**
     * 移除客户端连接
     */
    public void removeEmitter(String clientId) {
        SseEmitter emitter = emitters.remove(clientId);
        if (emitter != null) {
            emitter.complete();
        }
    }

    /**
     * 向特定客户端发送消息
     */
    public void sendMessageToClient(String clientId, String eventName, Object data) {
        SseEmitter emitter = emitters.get(clientId);
        if (emitter != null) {
            try {
                emitter.send(SseEmitter.event()
                    .name(eventName)
                    .data(data));
            } catch (IOException e) {
                emitters.remove(clientId);
                emitter.complete();
            }
        }
    }

    /**
     * 广播消息给所有连接的客户端
     */
    public void broadcastMessage(String eventName, Object data) {
        List<String> deadEmitters = new ArrayList<>();
        
        emitters.forEach((clientId, emitter) -> {
            try {
                emitter.send(SseEmitter.event()
                    .name(eventName)
                    .data(data));
            } catch (IOException e) {
                System.out.println("Client " + clientId + " disconnected");
                deadEmitters.add(clientId);
                emitter.complete();
            }
        });
        
        // 清理失效的连接
        deadEmitters.forEach(emitters::remove);
    }

    /**
     * 获取连接客户端数量
     */
    public int getConnectedClients() {
        return emitters.size();
    }
}
